/*
 * minimizer.cpp
 *
 *      Author: Jerome RD Soine and Christoph A Brand
 *      Institution: Schwarz goup, Institute for Theoretical Physics, Heidelberg University, Germany
 */

#include "../../../../../include/base/cell_models/ACM/numerics/minimizer.h"
#include "../../../../../include/base/cell_models/ACM/geometry/base/vertex.h"
#include "../../../../../include/base/cell_models/ACM/geometry/geomobject.h"
#include "../../../../../include/base/cell_models/ACM/constraints/constraint.h"

#include <vector>
#include <iostream>
#include <cmath>
#include <cstdlib>
#include "../../../../../include/base/cell_models/ACM/helper/helper.h"

using namespace std;

void Minimizer::iterate() {
	switch (mode) {
	case MODE_STEEPEST_DESCENT:
		iterateSteepestDescent();
		break;
	case MODE_CONJUGATE_GRADIENTS:
		iterateConjugateGradients();
		break;
	default:
		cout << "Mode <" << mode << "> not recognized. No iteration made."
				<< endl;
		break;
	}

}

void Minimizer::iterateConjugateGradients() {
	vector<double> forces, derivatives;
	initializeIterationStep();

	forces = derivatives = getTotalGradient();
	for (unsigned int i = 0; i < forces.size(); i++) {
		forces[i] *= -1.0;
	}

	projectForceOnConstraints(forces);

	if(firstConjugatedGradientStep){
		firstConjugatedGradientStep = false;
		conjugatedGradient = forces;
		lastGradient = forces;
	} else {
		double gammaZ = 0, gammaN=0, gamma=0;
		for (unsigned int i = 0; i < forces.size(); i++) {
			gammaZ += (forces[i]-lastGradient[i]) * forces[i];
			gammaN += lastGradient[i]*lastGradient[i];
		}

		gamma = gammaZ/gammaN;

		for (unsigned int i = 0; i < forces.size(); i++) {
			conjugatedGradient[i] = forces[i] + gamma * conjugatedGradient[i];
		}
		lastGradient=forces;
	}


	//	cout << "Find minimum" << endl;
	forces = conjugatedGradient;
	double min_displacement = goldenSectionSearch(forces, 0.001);
	for (unsigned int i = 0; i < forces.size(); i++) {
		forces[i] *= min_displacement;
	}

	//	cout << "Move vertices" << endl;
	g->moveVerticesBy(forces);

	vector<double> restoringMotion;
	restoringMotion = getConstraintRestoringMotion();
	g->moveVerticesBy(restoringMotion);


	//	cout << "Finished iteration" << endl << endl;
}

void Minimizer::iterateSteepestDescent() {
	//	cout << "Begin iteration" << endl;
	vector<double> forces, derivatives;
	initializeIterationStep();

	//	cout << "Get gradients" << endl;
	forces = derivatives = getTotalGradient();
	for (unsigned int i = 0; i < forces.size(); i++) {
		forces[i] *= -1.0;
	}

	//	cout << "Calc restoring motion" << endl;
	vector<double> restoringMotion;
	restoringMotion = getConstraintRestoringMotion();

	//	cout << "Project forces on constraints" << endl;
	projectForceOnConstraints(forces);

	//	cout << "Find minimum" << endl;
	double min_displacement = goldenSectionSearch(forces, 0.001);
	for (unsigned int i = 0; i < forces.size(); i++) {
		forces[i] *= min_displacement;
	}

	//	cout << "Move vertices" << endl;
	g->moveVerticesBy(forces);
	g->moveVerticesBy(restoringMotion);

	//	cout << "Finished iteration" << endl << endl;
}


void Minimizer::initializeIterationStep(){
	demandNewEnergyCalculation();
	demandNewGradientCalculation();
}

void Minimizer::demandNewEnergyCalculation(){
	for(unsigned int i=0; i<energies.size(); i++){
		energies[i]->setCalculateEnergies(true);
	}
}
void Minimizer::demandNewGradientCalculation(){
	for(unsigned int i=0; i<energies.size(); i++){
		energies[i]->setCalculateGradients(true);
	}
}

vector<double> Minimizer::getConstraintRestoringMotion() {
	if(constraints.size() == 0){
		return vector<double>(g->getVertices()->size()*3,0.0);
	} else if (constraints.size() == 1) {
		return getConstraintRestoringMotionOfOneConstraint();
	} else if (constraints.size() == 2) {
		return getConstraintRestoringMotionOfTwoConstraints();
	} else {
		cout << "Error: Cannot handle more than two constraint at the moment!"
				<< endl;
		exit(1);
	}
}

vector<double> Minimizer::getConstraintRestoringMotionOfOneConstraint() {
	vector<Vertex*> vertices = *g->getVertices();
	vector<double> restoringMotion;
	restoringMotion.resize(vertices.size() * 3);

	double deviation = constraints[0]->getDeviation();

	Vertex gradients[vertices.size()];
	double absoluteGradient = 0;
	for (unsigned int i = 0; i < vertices.size(); i++) {
		Vertex gradient(0, 0, 0);
		if (!vertices[i]->isFixed()) {
			gradient = constraints[0]->getGradient(*vertices[i]);
		}
		gradients[i] = gradient;
		absoluteGradient += gradient * gradient;

	}

	double scaleFactor = (-1) * deviation / absoluteGradient;

	for (unsigned int i = 0; i < vertices.size(); i++) {
		restoringMotion[3 * i + 0] += scaleFactor * gradients[i].getX();
		restoringMotion[3 * i + 1] += scaleFactor * gradients[i].getY();
		restoringMotion[3 * i + 2] += scaleFactor * gradients[i].getZ();
	}
	return restoringMotion;
}

vector<double> Minimizer::getConstraintRestoringMotionOfTwoConstraints() {
	vector<double> deviations;
	for (unsigned int i = 0; i < constraints.size(); i++) {
		deviations.push_back(constraints[i]->getDeviation());
	}

	vector<double> A;
	A.push_back(
			Constraint::getSumOfScalarProductsOfGradientsOverAllVertices(
					constraints[0], constraints[0]));
	A.push_back(
			Constraint::getSumOfScalarProductsOfGradientsOverAllVertices(
					constraints[1], constraints[0]));
	A.push_back(
			Constraint::getSumOfScalarProductsOfGradientsOverAllVertices(
					constraints[0], constraints[1]));
	A.push_back(
			Constraint::getSumOfScalarProductsOfGradientsOverAllVertices(
					constraints[1], constraints[1]));
	Helper::invert2x2Matrix(A);
	vector<double> coefficients;
	coefficients.push_back(
			(-1) * deviations[0] * A[0] + (-1) * deviations[1] * A[1]);
	coefficients.push_back(
			(-1) * deviations[0] * A[2] + (-1) * deviations[1] * A[3]);

	vector<double> restore;
	vector<Vertex*> vertices = *g->getVertices();
	restore.reserve(vertices.size() * 3);
	for (unsigned int i = 0; i < vertices.size(); i++) {
		Vertex restoringMotion(0, 0, 0);
		if (!vertices[i]->isFixed()) {
			for (unsigned int j = 0; j < constraints.size(); j++) {
				restoringMotion += constraints[j]->getGradient(*vertices[i])
						* coefficients[j];
			}
		}
		restore.push_back(restoringMotion.getX());
		restore.push_back(restoringMotion.getY());
		restore.push_back(restoringMotion.getZ());
	}
	return restore;
}

void Minimizer::projectForceOnConstraints(vector<double> & force) {
	if(constraints.size() == 0){
		// do nothing.
	} else	if (constraints.size() == 1) {
		projectForceOnOneConstraint(force);
	} else if (constraints.size() == 2) {
		projectForceOnTwoConstraints(force);
	} else {
		cout << "Error: Cannot handle more than two constraints at the moment!"
				<< endl;
		exit(1);
	}
}

void Minimizer::projectForceOnOneConstraint(vector<double> & forces) {
	vector<Vertex*> vertices = *g->getVertices();

	Vertex gradients[vertices.size()];
	double sumOfForceTimesGradient = 0, sumOfGradientTimesGradient = 0;
	for (unsigned int i = 0; i < vertices.size(); i++) {
		Vertex force(forces[3 * i + 0], forces[3 * i + 1], forces[3 * i + 2]);
		Vertex gradient(0,0,0);
		if(!vertices[i]->isFixed()) gradient= constraints[0]->getGradient(*vertices[i]);
		gradients[i] = gradient;
		sumOfForceTimesGradient += force * gradient;
		sumOfGradientTimesGradient += gradient * gradient;
	}
	double scale = sumOfForceTimesGradient / sumOfGradientTimesGradient;

	for (unsigned int i = 0; i < vertices.size(); i++) {
		Vertex diff = gradients[i] * scale;
		forces[3 * i + 0] -= diff.getX();
		forces[3 * i + 1] -= diff.getY();
		forces[3 * i + 2] -= diff.getZ();
	}
}
void Minimizer::projectForceOnTwoConstraints(vector<double> & forces) {
	vector<double> forceGradientSums;
	for (unsigned int i = 0; i < constraints.size(); i++) {
		forceGradientSums.push_back(
				constraints[i]->getSumOfScalarProductsOfGradientsAndForcesOverAllVertices(
						forces));
	}
	vector<double> A;
	A.push_back(
			Constraint::getSumOfScalarProductsOfGradientsOverAllVertices(
					constraints[0], constraints[0]));
	A.push_back(
			Constraint::getSumOfScalarProductsOfGradientsOverAllVertices(
					constraints[1], constraints[0]));
	A.push_back(
			Constraint::getSumOfScalarProductsOfGradientsOverAllVertices(
					constraints[0], constraints[1]));
	A.push_back(
			Constraint::getSumOfScalarProductsOfGradientsOverAllVertices(
					constraints[1], constraints[1]));
	Helper::invert2x2Matrix(A);
	vector<double> coefficients;
	coefficients.push_back(
			forceGradientSums[0] * A[0] + forceGradientSums[1] * A[1]);
	coefficients.push_back(
			forceGradientSums[0] * A[2] + forceGradientSums[1] * A[3]);

	vector<Vertex*> vertices = *g->getVertices();
	for (unsigned int i = 0; i < vertices.size(); i++) {
		Vertex diff(0, 0, 0);
		if(!vertices[i]->isFixed()){
			for (unsigned int j = 0; j < constraints.size(); j++) {
				diff += constraints[j]->getGradient(*vertices[i]) * coefficients[j];
			}
		}
		forces[3 * i + 0] -= diff.getX();
		forces[3 * i + 1] -= diff.getY();
		forces[3 * i + 2] -= diff.getZ();
	}
}

double Minimizer::getTotalEnergy() {
	double energy = 0;
	for (unsigned int i = 0; i < energies.size(); i++) {
		energy += energies[i]->getEnergy();
	}
	return energy;
}

vector<double> Minimizer::getTotalGradient() {
	vector<double> tot_grad;
	vector<Vertex*> vertices = *g->getVertices();
	tot_grad.reserve(vertices.size() * 3);
	for (unsigned int i = 0; i < vertices.size(); i++) {
		Vertex grad = getTotalGradientOfOneVertex(*vertices[i]);
		tot_grad.push_back(grad.getX());
		tot_grad.push_back(grad.getY());
		tot_grad.push_back(grad.getZ());
	}
	return tot_grad;
}

Vertex Minimizer::getTotalGradientOfFixedPoints(Vertex & v)
{
	Vertex ret(0, 0, 0);
	if (!v.isFixed())
		return ret;
	else {
		for (unsigned int i = 0; i < energies.size(); i++) {
			ret += energies[i]->getGradient(v);
		}
		return ret;
	}
}

Vertex Minimizer::getTotalGradientOfOneVertex(Vertex & v) {
	Vertex ret(0, 0, 0);
	if (v.isFixed())
		return ret;
	else {
		for (unsigned int i = 0; i < energies.size(); i++) {
			ret += energies[i]->getGradient(v);
		}
		return ret;
	}
}

double Minimizer::goldenSectionSearch(vector<double> & searchDirection,
		double tol) {
	double min_displacement, after_min_displacement;
	//listForces(searchDirection);
	bracketMinimum(min_displacement, after_min_displacement, searchDirection); // Minimum is bracketed in (0, min_displacement, after_min_displacement)
	const double R = 0.61803399, C = 1.0 - R, zero_displacement = 0.0,
			max_displacement = after_min_displacement;
	double x0 = zero_displacement, x1, x2, x3 = max_displacement;
	if (abs(min_displacement)
			> abs(after_min_displacement - min_displacement)) {
		x1 = C * min_displacement;
		x2 = min_displacement;
	} else {
		x1 = min_displacement;
		x2 = min_displacement + R * min_displacement;
	}

	double e1 = oneDimensionalEnergyEval(searchDirection, x1), e2 =
			oneDimensionalEnergyEval(searchDirection, x2);

	int i=0;

	while (abs(x3 - x0) > tol * (abs(x1) + abs(x2))) {
		if (e2 < e1) {
			x0 = x1;
			x1 = x2;
			x2 = R * x2 + C * x3;
			e1 = e2;
			e2 = oneDimensionalEnergyEval(searchDirection, x2);
		} else {
			x3 = x2;
			x2 = x1;
			x1 = R * x1 + C * x0;
			e2 = e1;
			e1 = oneDimensionalEnergyEval(searchDirection, x1);
		}
		if(i++>100) {
			throw ERROR_TOO_MANY_ITERATIONS_IN_GOLDEN_SECTION_SEARCH;
		}
	}

	if (e1 < e2) {
		//cout << "Minimum at displacement " << x1 << endl;
		return x1;
	} else {
		//cout << "Minimum at displacement " << x2 << endl;
		return x2;
	}

}

void Minimizer::bracketMinimum(double & min_displacement,
		double & after_min_displacement, vector<double> & searchDirection) {
	double eA, eB, eC;

	const double gold = 1.618034;
	double displacement = 0.01;

	eA = getTotalEnergy();
	eC = oneDimensionalEnergyEval(searchDirection, displacement);
	int i = 0;
	if (eC > eA) {
		while (eC > eA && i < 10) {
			displacement *= 0.1;
			eC = oneDimensionalEnergyEval(searchDirection, displacement);
			i++;
		}
	}
	//listForces(searchDirection);
	if (eC > eA) {
		throw ERROR_BRACKETING_MINIMUM;
	}
	after_min_displacement = displacement;
	i = 0;
	do {
		i++;
		eB = eC;
		min_displacement = after_min_displacement;
		displacement *= gold;
		after_min_displacement += displacement;
		eC = oneDimensionalEnergyEval(searchDirection, after_min_displacement);
	} while (eC < eB);
}

double Minimizer::oneDimensionalEnergyEval(vector<double> & direction,
		double displacement_scalar) {
	vector<double> displacement = direction;
	for (unsigned int i = 0; i < direction.size(); i++) {
		displacement[i] *= displacement_scalar;
	}
	g->moveVerticesBy(displacement);
	demandNewEnergyCalculation();
	double e = getTotalEnergy();
	g->moveVerticesBackBy(displacement);
	demandNewEnergyCalculation();
	return e;
}

void Minimizer::listForces(vector<double> & forces) {
	Vertex highForce(0,0,0), pos(0,0,0);
	vector<Vertex*> vertices = *g->getVertices();
	for (unsigned int i = 0; i < vertices.size(); i++) {
		Vertex force(0, 0, 0);
		force.setX(forces[3 * i + 0]);
		force.setY(forces[3 * i + 1]);
		force.setZ(forces[3 * i + 2]);
		if(force.abs() > highForce.abs()){
			highForce = force;
			pos = *vertices[i];
		}
		//cout << "Vertex " << i << ": " << (*vertices)[i]->toString()	<< " - Force: " << force.toString() << endl;
	}
	cout << endl << "Highest absolute force: ";
	cout << highForce.abs() << " at Vertex " << pos.toString() << " in direction " << (highForce/highForce.abs()).toString()<< endl;

}

void Minimizer::printEnergyLine(vector<double> direction, double displacement,
		int iterations) {
	for (int i = 0; i < iterations; i++) {
		std::cout << "Displacement: " << i * displacement << " - Energy: "
				<< oneDimensionalEnergyEval(direction, i * displacement)
				<< std::endl;
	}
}

